{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Th91ofnQWr8Y"
      },
      "source": [
        "## Dataset building + XTTS finetuning and inference\n",
        "\n",
        "#### Running the demo\n",
        "To start the demo run the first two cells (ignore pip install errors in the first one)\n",
        "\n",
        "Then click on the link `Running on public URL: ` when the demo is ready.\n",
        "\n",
        "#### Downloading the results\n",
        "\n",
        "You can run cell [3] to zip and download default dataset path\n",
        "\n",
        "You can run cell [4] to zip and download the latest model you trained"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cdWKA_xFqkKq"
      },
      "source": [
        "### Installing the requirements"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lmUUQqdN6BXk"
      },
      "outputs": [],
      "source": [
        "!rm -rf coqui-ai-TTS/ # delete repo to be able to reinstall if needed\n",
        "!git clone -q https://github.com/idiap/coqui-ai-TTS.git\n",
        "!pip install -q -e coqui-ai-TTS\n",
        "!pip install -q gradio faster_whisper"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g7rNt1e2qtDP"
      },
      "source": [
        "### Running the gradio UI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zd2xo_7a8wyj"
      },
      "outputs": [],
      "source": [
        "!python -m TTS.demos.xtts_ft_demo.xtts_demo --batch_size 2 --num_epochs 6"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oXEBRA_kq23i"
      },
      "source": [
        "### Downloading the dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dBxgdKcvi4kO"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "!zip -q -r dataset.zip /tmp/xtts_ft/dataset\n",
        "files.download('dataset.zip')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZKzoP53Nq_rJ"
      },
      "source": [
        "### Downloading the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NpfdzHvKaX8M"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "import os\n",
        "import glob\n",
        "import torch\n",
        "\n",
        "def find_latest_best_model(folder_path):\n",
        "    search_path = os.path.join(folder_path, '**', 'best_model.pth')\n",
        "    files = glob.glob(search_path, recursive=True)\n",
        "    latest_file = max(files, key=os.path.getctime, default=None)\n",
        "    return latest_file\n",
        "\n",
        "model_path = find_latest_best_model(\"/tmp/xtts_ft/run/training/\")\n",
        "checkpoint = torch.load(model_path, map_location=torch.device(\"cpu\"))\n",
        "del checkpoint[\"optimizer\"]\n",
        "for key in list(checkpoint[\"model\"].keys()):\n",
        "    if \"dvae\" in key:\n",
        "        del checkpoint[\"model\"][key]\n",
        "torch.save(checkpoint, \"model.pth\")\n",
        "model_dir = os.path.dirname(model_path)\n",
        "files.download(os.path.join(model_dir, 'config.json'))\n",
        "files.download(os.path.join(model_dir, 'vocab.json'))\n",
        "files.download('model.pth')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eh9_SusYdRE4"
      },
      "source": [
        "### Copy files to your google drive\n",
        "\n",
        "The two previous cells are a requirement for this step but it can be much faster"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "piLAaVHSdQs5"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "import shutil\n",
        "drive.mount('/content/drive')\n",
        "!mkdir /content/drive/MyDrive/XTTS_ft_colab\n",
        "shutil.copy(os.path.join(model_dir, 'config.json'), \"/content/drive/MyDrive/XTTS_ft_colab/config.json\")\n",
        "shutil.copy(os.path.join(model_dir, 'vocab.json'), \"/content/drive/MyDrive/XTTS_ft_colab/vocab.json'\")\n",
        "shutil.copy('model.pth', \"/content/drive/MyDrive/XTTS_ft_colab/model.pth\")\n",
        "shutil.copy('dataset.zip', \"/content/drive/MyDrive/XTTS_ft_colab/dataset.zip\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
